# -*- coding:utf-8 -*-

from fastapi import APIRouter, Depends, Form
from sqlalchemy.ext.asyncio import AsyncSession
from starlette.requests import Request
from typing import List
from config.enums import BusinessType
from config.get_db import get_db
from module_admin.entity.vo.import_vo import ImportModel
from module_admin.service.import_service import ImportService
from module_admin.service.login_service import LoginService
from module_admin.aspect.data_scope import GetDataScope
from module_admin.aspect.interface_auth import CheckUserInterfaceAuth
from module_admin.entity.vo.user_vo import CurrentUserModel
from module_admin.annotation.log_annotation import Log
from module_stock.entity.do.dc_index_do import DcIndex
from module_stock.entity.do.dc_member_do import DcMember
from utils.response_util import ResponseUtil
from utils.common_util import bytes2file_response

from module_stock.entity.vo.stock_basic_vo import StockBasicPageModel, StockBasicModel
from module_stock.service.stock_basic_service import StockBasicService
from module_stock.service.daily_service import DailyService
stockBasicController = APIRouter(prefix='/stock/basic', dependencies=[Depends(LoginService.get_current_user)])
import pandas as pd
from module_stock.entity.vo.daily_vo import DailyPageModel
import math
from sqlalchemy import select, and_, func
from sqlalchemy.orm import aliased

@stockBasicController.get('/list', dependencies=[Depends(CheckUserInterfaceAuth('stock:basic:list'))])
async def get_stock_basic_list(
        request: Request,
        query_db: AsyncSession = Depends(get_db),
        page_query: StockBasicPageModel = Depends(StockBasicPageModel.as_query),
        data_scope_sql: str = Depends(GetDataScope('StockBasic'))
):

    stock_basic_result = await StockBasicService.get_stock_basic_list_nopage(query_db, page_query, data_scope_sql)
    stock_basic_df = pd.DataFrame(stock_basic_result)

    # 解析 otherInfo 字段，提取涨跌幅数据
    # stock_basic_df = await parse_other_info_data(stock_basic_df)
    # 拼接daily数据
    merged_daily_df = await mergeDaily(query_db, stock_basic_df)

    # 拼接东方财富板块数据
    merged_dc_index_df = await mergeDcIndexData(query_db, merged_daily_df)

    # 转换回原来的格式
    merged_data = merged_dc_index_df.to_dict('records')

    # 清理可能存在的非法浮点值（inf, nan等）
    # def clean_data(obj):
    #     if isinstance(obj, dict):
    #         return {k: clean_data(v) for k, v in obj.items()}
    #     elif isinstance(obj, list):
    #         return [clean_data(item) for item in obj]
    #     elif isinstance(obj, float):
    #         if math.isnan(obj) or math.isinf(obj):
    #             return None
    #     return obj
    #
    # cleaned_data = clean_data(merged_data)

    return ResponseUtil.success(data=merged_data)


async def parse_other_info_data(stock_basic_df):
    """
    解析 stock_basic_df 中的 otherInfo 字段，提取涨跌幅数据作为新字段
    Args:
        stock_basic_df: 包含 otherInfo 字段的 DataFrame
    Returns:
        DataFrame: 添加了涨跌幅字段的 DataFrame
    """
    # 解析 otherInfo 字段，提取涨跌幅数据
    if 'otherInfo' in stock_basic_df.columns:
        import json

        def extract_price_changes(row):
            if pd.notna(row['otherInfo']) and row['otherInfo']:
                try:
                    other_info = json.loads(row['otherInfo'])
                    price_changes = other_info.get('price_changes', {})

                    # 提取各种复权方式的涨跌幅
                    for method in ['raw', 'qfq', 'hfq']:
                        changes = price_changes.get(method, {})
                        for period, value in changes.items():
                            # 为不同复权方式创建不同的列名
                            if method == 'raw':
                                col_name = period  # 原始数据直接使用period名称，如 change_3d
                            else:
                                col_name = f"{method}_{period}"  # 复权数据加上前缀，如 qfq_change_3d
                            row[col_name] = value
                except (json.JSONDecodeError, KeyError, TypeError):
                    pass
            return row

        # 应用解析函数到每一行
        stock_basic_df = stock_basic_df.apply(extract_price_changes, axis=1)

    return stock_basic_df


async def mergeDaily(query_db, stock_basic_df):
    latest_daily_query = DailyPageModel(
        page_num=0,
        page_size=-1,  # 设置一个较大的数值以获取所有数据
        trade_date=None
    )
    # 先获取最新的交易日期
    latest_trade_date = await DailyService.get_latest_trade_date(query_db)
    # 如果存在最新交易日期，则使用它进行查询
    if latest_trade_date:
        latest_daily_query.trade_date = latest_trade_date

        daily_result = await DailyService.get_daily_list_nopage(query_db, latest_daily_query)
        daily_df = pd.DataFrame(daily_result)
        # 使用merge进行数据合并
        if not daily_df.empty:
            # 合并数据
            merged_df = stock_basic_df.merge(daily_df, on='tsCode', how='left', suffixes=('', '_daily'))

            # 找出所有带_daily后缀的重复列并删除
            daily_columns = [col for col in merged_df.columns if col.endswith('_daily')]
            merged_df = merged_df.drop(columns=daily_columns)
            # 直接清理merged_df中的无效浮点值
            merged_df = merged_df.map(
                lambda x: None if isinstance(x, float) and (math.isnan(x) or math.isinf(x)) else x)
        else:
            # 如果没有daily数据，直接使用stock_basic数据
            merged_df = stock_basic_df
    else:
        merged_df = stock_basic_df
    return merged_df


async def mergeDcIndexData(query_db, merged_daily_df):
    """
    在数据库层面进行复杂关联查询，获取成分股及其所有关联板块信息，最后与merged_daily_df合并数据
    """
    # 获取最新的交易日期
    latest_trade_date_stmt = select(func.max(DcMember.trade_date))
    latest_trade_date_result = await query_db.execute(latest_trade_date_stmt)
    latest_trade_date = latest_trade_date_result.scalar()

    if not latest_trade_date:
        return merged_daily_df

    # 复杂关联查询：一次查询获取成分股及其所有关联板块信息
    complex_query = select(
        DcMember.con_code,
        DcMember.name.label('member_name'),
        DcIndex.ts_code.label('index_code'),
        DcIndex.name.label('index_name'),
        DcIndex.pct_change.label('index_pct_change'),
        DcIndex.up_num.label('index_up_num'),
        DcIndex.down_num.label('index_down_num'),
        DcIndex.leading.label('index_leading'),
        DcIndex.leading_code.label('index_leading_code'),
        DcIndex.leading_pct.label('index_leading_pct'),
        DcIndex.total_mv.label('index_total_mv'),
        DcIndex.turnover_rate.label('index_turnover_rate')
    ).select_from(
        DcMember.__table__.join(
            DcIndex.__table__,
            DcMember.ts_code == DcIndex.ts_code
        )
    ).where(
        and_(
            DcMember.trade_date == latest_trade_date,
            DcIndex.trade_date == latest_trade_date
        )
    ).order_by(
        DcMember.con_code, DcIndex.ts_code
    )

    complex_result = await query_db.execute(complex_query)
    complex_data = complex_result.all()

    # 处理查询结果，按con_code分组构建板块信息
    con_code_to_indices = {}
    for row in complex_data:
        con_code = row.con_code
        if con_code not in con_code_to_indices:
            con_code_to_indices[con_code] = []

        index_info = {
            'index_code': row.index_code,
            'index_name': row.index_name,
            'index_pct_change': float(row.index_pct_change) if row.index_pct_change is not None else None,
            'index_up_num': row.index_up_num,
            'index_down_num': row.index_down_num,
            'index_leading': row.index_leading,
            'index_leading_code': row.index_leading_code,
            'index_leading_pct': float(row.index_leading_pct) if row.index_leading_pct is not None else None,
            'index_total_mv': float(row.index_total_mv) if row.index_total_mv is not None else None,
            'index_turnover_rate': float(row.index_turnover_rate) if row.index_turnover_rate is not None else None
        }
        con_code_to_indices[con_code].append(index_info)

    # 将板块信息合并到merged_daily_df中
    def add_index_info(row):
        ts_code = row['tsCode']
        # 注意：这里假设tsCode对应的是con_code
        if ts_code in con_code_to_indices:
            row['index_info'] = con_code_to_indices[ts_code]
        else:
            row['index_info'] = []
        return row

    # 应用板块信息到每一行数据
    merged_daily_df = merged_daily_df.apply(add_index_info, axis=1)

    return merged_daily_df


@stockBasicController.get('/getById/{stockBasicId}', dependencies=[Depends(CheckUserInterfaceAuth('stock:basic:list'))])
async def get_stock_basic_by_id(
        request: Request,
        stockBasicId: int,
        query_db: AsyncSession = Depends(get_db),
        data_scope_sql: str = Depends(GetDataScope('StockBasic'))
):
    stock_basic = await StockBasicService.get_stock_basic_by_id(query_db, stockBasicId)
    return ResponseUtil.success(data=stock_basic)


@stockBasicController.post('/add', dependencies=[Depends(CheckUserInterfaceAuth('stock:basic:add'))])
@Log(title='stock_basic', business_type=BusinessType.INSERT)
async def add_stock_basic (
    request: Request,
    add_model: StockBasicModel,
    query_db: AsyncSession = Depends(get_db),
    current_user: CurrentUserModel = Depends(LoginService.get_current_user),
):

    add_model.create_by = current_user.user.user_id
    add_model.dept_id = current_user.user.dept_id
    add_dict_type_result = await StockBasicService.add_stock_basic(query_db, add_model)
    return ResponseUtil.success(data=add_dict_type_result)

@stockBasicController.put('/update', dependencies=[Depends(CheckUserInterfaceAuth('stock:basic:edit'))])
@Log(title='stock_basic', business_type=BusinessType.UPDATE)
async def update_stock_basic(
    request: Request,
    edit_model: StockBasicModel,
    query_db: AsyncSession = Depends(get_db),
    current_user: CurrentUserModel = Depends(LoginService.get_current_user),
):
    add_dict_type_result = await StockBasicService.update_stock_basic(query_db, edit_model)
    return ResponseUtil.success(data=add_dict_type_result)


@stockBasicController.delete('/delete/{stockBasicIds}', dependencies=[Depends(CheckUserInterfaceAuth('stock:basic:del'))])
@Log(title='stock_basic', business_type=BusinessType.DELETE)
async def del_stock_basic(
    request: Request,
    stockBasicIds: str,
    query_db: AsyncSession = Depends(get_db),
    current_user: CurrentUserModel = Depends(LoginService.get_current_user),
):
    ids = stockBasicIds.split(',')
    del_result = await StockBasicService.del_stock_basic(query_db, ids)
    return ResponseUtil.success(data=del_result)

@stockBasicController.post('/export', dependencies=[Depends(CheckUserInterfaceAuth('stock:basic:export'))])
@Log(title='stock_basic', business_type=BusinessType.EXPORT)
async def export_stock_basic(
    request: Request,
    stock_basic_form: StockBasicPageModel = Form(),
    query_db: AsyncSession = Depends(get_db),
    data_scope_sql: str = Depends(GetDataScope('StockBasic')),
):
    # 获取全量数据
    export_result = await StockBasicService.export_stock_basic_list(
        query_db, stock_basic_form, data_scope_sql
    )
    return ResponseUtil.streaming(data=bytes2file_response(export_result))

@stockBasicController.post('/import', dependencies=[Depends(CheckUserInterfaceAuth('stock:basic:import'))])
async def import_stock_basic(request: Request,
                      import_model: ImportModel,
                      query_db: AsyncSession = Depends(get_db),
                      current_user: CurrentUserModel = Depends(LoginService.get_current_user)
    ):
    """
    导入数据
    """
    await ImportService.import_data(query_db, import_model, current_user)
    return ResponseUtil.success()